import zmq
import time
import signal
import sys
import socket


class ZeroMQClient:
    def __init__(self, mq_url, task_type, worker):
        # 获取worker的类名
        self.worker_type = worker.__class__.__name__ if worker else ''

        # 对应zmq的属性
        self.mq_url = mq_url
        self.task_type = task_type
        self.worker = worker
        self.context = None
        self.socket = None
        self.running = True

        # 设置信号处理
        # signal.signal(signal.SIGINT, self.signal_handler)
        # signal.signal(signal.SIGTERM, self.signal_handler)

        # 发送心跳信息
        self.heartbeat_interval = 5
        self.custom_info = {}

        self.ip = self.get_local_ip()
        print(f'本地ip地址: {self.ip}')

    def get_local_ip(self):
        """获取当前机器的 IP 地址"""
        try:
            # 创建一个 UDP 套接字
            s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            s.connect(("8.8.8.8", 80))  # 连接到公共 DNS 服务器
            ip = s.getsockname()[0]  # 获取本地 IP 地址
            s.close()
            return ip
        except Exception as e:
            print(f"获取 IP 地址失败: {e}")
            return "Unknown IP"

    def setCustomInfo(self, custom_info):
        self.custom_info = custom_info

    def signal_handler(self, signum, frame):
        print("接收到关闭信号，准备退出...")
        self.running = False
        self.cleanup()
        sys.exit(0)

    def start(self):
        self.context = zmq.Context()
        self.socket = self.context.socket(zmq.DEALER)
        self.socket.connect(self.mq_url)
        print("[zmq] ZeroMQ 套接字已连接...")

        last_heartbeat_time = time.time()
        try:
            while self.running:
                # print(f"Worker 正在等待 {self.task_type} 任务...")

                # 发送心跳消息
                if time.time() - last_heartbeat_time > self.heartbeat_interval:
                    self.socket.send_json({
                        "taskType": 'heartbeat',
                        "workerType": self.worker_type,
                        "action": "heartbeat",
                        "payload": {
                            "ip": self.ip,
                            "customInfo": self.custom_info
                        }
                    })
                    last_heartbeat_time = time.time()

                # 发送请求，获取任务
                self.socket.send_json(
                    {
                        "taskType": self.task_type,
                        "workerType": self.worker_type,
                        "action": "requestTask"
                    })
                # print(f"Worker 正在等待 {self.task_type} 任务...")

                # 接收任务
                message = self.socket.recv_json()
                task = message['task']

                if task:
                    # 处理消息
                    result = self.worker.process(task)

                    # 返回结果
                    self.socket.send_json({
                        "taskType": "task_result",
                        "workerType": self.worker_type,
                        "action": "submitResult",
                        "payload": result
                    })

                    # print(f"结果已发布到频道task_result: {result}")
                    print(f"结果已发布到频道task_result")

                time.sleep(1)
        except zmq.ZMQError as e:
            print(f"ZeroMQ 错误: {e}")
        except KeyboardInterrupt:
            print("接收到中断信号，准备退出...")
        finally:
            self.cleanup()

    def cleanup(self):
        """清理 ZeroMQ 套接字和上下文"""
        if self.socket:
            self.socket.close()
            print("[zmq] ZeroMQ 套接字已关闭...")
        # if self.context:
        #     self.context.term()
        #     print("ZeroMQ 上下文已关闭")
